#include "stdafx.h"
#include <algorithm>
#include <string>
#include "proc.h"

namespace xm
{
	imagex_t::imagex_t() : export_loads(0), export_ord_base(0), _is64(0)
	{
		memset( &_file_hdr, 0, sizeof(_file_hdr) );
		memset( &_nt_hdr64, 0, sizeof(_nt_hdr64) );
	}

	long imagex_t::open( IUnknown * file )
	{
		iptr<IKBStream> fo(file);
		if( !fo ) return E_INVALIDARG;

		IMAGE_DOS_HEADER dh = {};
		
		size_t eff = 0;
		RFAILED( fo->Readp( 0, &dh, sizeof(dh), &eff ) );
		if( eff != sizeof(dh) ) return E_FAIL;
		
		if( dh.e_magic != IMAGE_DOS_SIGNATURE )
			return E_FAIL;
		
		uint32_t sign = 0;
		RFAILED( fo->Readp( dh.e_lfanew, &sign, sizeof(sign), &eff ) );
		if( eff != sizeof(sign) ) return E_FAIL;

		if( sign != IMAGE_NT_SIGNATURE ) return E_FAIL;

		RFAILED( fo->Read( &_file_hdr, sizeof(_file_hdr), &eff ) );
		if( eff != sizeof(_file_hdr) ) 
			return E_FAIL;

		// 386/AMD64
		if( _file_hdr.Machine != IMAGE_FILE_MACHINE_I386 &&
			_file_hdr.Machine != IMAGE_FILE_MACHINE_AMD64 )
			return E_FAIL;

		_is64 = (_file_hdr.Machine == IMAGE_FILE_MACHINE_AMD64);

		size_t size_of_opt_std = _is64 ? sizeof(IMAGE_OPTIONAL_HEADER64) : sizeof(IMAGE_OPTIONAL_HEADER32);
		size_t size_of_opt = _file_hdr.SizeOfOptionalHeader;

		size_of_opt = size_of_opt < size_of_opt_std ? size_of_opt : size_of_opt_std;
		if( _is64 )
		{
			RFAILED( fo->Read( &_nt_hdr64, size_of_opt, &eff ) );
		}
		else
		{
			RFAILED( fo->Read( &_nt_hdr32, size_of_opt, &eff ) );
		}

		if( eff != size_of_opt ) return E_FAIL;

		if( _is64 != (_nt_hdr64.Magic == IMAGE_NT_OPTIONAL_HDR64_MAGIC) ) return E_FAIL;
		if( _is64 == (_nt_hdr32.Magic == IMAGE_NT_OPTIONAL_HDR32_MAGIC) ) return E_FAIL;

		//size_t fil_align = _is64 ? _nt_hdr64.FileAlignment : _nt_hdr32.FileAlignment;
		//if( fil_align != 0x200 ) fil_align = 0x200;
		//size_t img_align = _is64 ? _nt_hdr64.SectionAlignment : _nt_hdr32.SectionAlignment;
		//if( img_align < fil_align ) img_align = fil_align;

		size_t fil_align = 0x200, img_align = VM_PAGE_SIZE;

		size_t image_size = 0;
		if( _is64 )
			image_size = (_nt_hdr64.SizeOfImage = __align_big(_nt_hdr64.SizeOfImage, img_align ));
		else
			image_size = (_nt_hdr32.SizeOfImage = __align_big(_nt_hdr32.SizeOfImage, img_align ));

		size_t nSections = _file_hdr.NumberOfSections + 1;
		_sections.resize( nSections );
	
		size_t size_of_sectab = sizeof(IMAGE_SECTION_HEADER)*_file_hdr.NumberOfSections;
		RFAILED( fo->Read( &_sections[1], size_of_sectab, &eff ) );

		if( size_of_sectab != eff ) 
			return E_FAIL;

		uint32_t rva_min = -1, rva_max = 0;
		
		DWORD prev_ve = 0;
		for( size_t i = 1; i < nSections; ++ i )
		{
			DWORD& this_va = _sections[i].VirtualAddress;

			if( prev_ve && prev_ve > this_va ) 
				return E_FAIL;
			if( this_va % img_align ) 
				return E_FAIL;

			DWORD& this_vs = _sections[i].Misc.VirtualSize;
			if( !this_vs ) this_vs = _sections[i].SizeOfRawData;
			this_vs = __align_big(this_vs, img_align );
			
			DWORD this_ve = this_va + this_vs;
			if( this_va < rva_min ) rva_min = this_va;
			if( this_ve > rva_max ) rva_max = this_ve;
			
			prev_ve = this_ve;
		}

		if( !rva_min ) return E_FAIL;
		if( rva_max > image_size ) return E_FAIL;

		// fix me 
		size_t header_size = dh.e_lfanew + sizeof(uint32_t) + sizeof(_file_hdr) + size_of_opt + size_of_sectab;
		if( header_size > rva_min ) header_size = rva_min;
		IMAGE_SECTION_HEADER header = {'{','h','e','a','d','e','r','}' };
		header.Characteristics = IMAGE_SCN_MEM_READ;
		header.PointerToRawData = 0;
		header.SizeOfRawData = header_size;
		header.VirtualAddress = 0;
		header.Misc.VirtualSize = rva_min;
		_sections[0] = header;

		//////////////////////////////////////////////////////////////////////////
		//
		// LOAD TLS
		//
		//IMAGE_DATA_DIRECTORY * dde = get_dde(IMAGE_DIRECTORY_ENTRY_TLS);
		//do {

		//	if( !dde || !dde->VirtualAddress ) break;
		//	
		//	size_t want = _is64 ? sizeof(_tls64) : sizeof(_tls32);
		//	if( FAILED( Readp( dde->VirtualAddress, &_tls64, want, &eff ) ) )
		//		break;

		//}while(0);

		_file = fo;
		return S_OK;
	}

	uint64_t imagex_t::image_base()
	{
		if( !_file ) return 0;
		return _is64 ? _nt_hdr64.ImageBase : _nt_hdr32.ImageBase;
	}
	size_t imagex_t::entry_point()
	{
		return _is64 ? _nt_hdr64.AddressOfEntryPoint : _nt_hdr32.AddressOfEntryPoint;
	}
	size_t	imagex_t::num_of_dde()
	{
		return _is64 ? _nt_hdr64.NumberOfRvaAndSizes : _nt_hdr32.NumberOfRvaAndSizes;
	}
	IMAGE_DATA_DIRECTORY * imagex_t::get_dde( )
	{
		return _is64 ? _nt_hdr64.DataDirectory : _nt_hdr32.DataDirectory;
	}
	IMAGE_DATA_DIRECTORY * imagex_t::get_dde( size_t i )
	{
		if( i >= num_of_dde() ) return NULL;
		return get_dde() + i;
	}

	size_t imagex_t::relocatable()
	{
		IMAGE_DATA_DIRECTORY * dde = get_dde( IMAGE_DIRECTORY_ENTRY_BASERELOC );
		if( !dde ) return 0;
		if( !dde->VirtualAddress || !dde->Size ) 
			return 0;
		return 1;
	}

	IMAGE_SECTION_HEADER * imagex_t::get_section( size_t i )
	{
		if( i >= _sections.size() )
			return NULL;
		return &_sections[i];
	}

	IMAGE_SECTION_HEADER * imagex_t::find_section( uint32_t rva, size_t * pi )
	{
		size_t scn_num = _sections.size();
		for( size_t i = 0; i < scn_num; ++ i )
		{
			size_t this_va = _sections[i].VirtualAddress;
			if( this_va > rva ) return NULL;
			size_t this_ve = _sections[i].VirtualAddress + _sections[i].Misc.VirtualSize;
			if( rva < this_va ) continue;
			if( rva >= this_ve ) continue;
			if( pi ) *pi = i;
			return &_sections[i];
		}
		return NULL;
	}
	STDMETHODIMP_(BOOL) imagex_t::IsVirtual()
	{ 
		return _file->IsVirtual(); 
	}
	STDMETHODIMP imagex_t::Readp( uint64_t rva, void * pbuf, size_t want, size_t * eff )
	{
		if( eff ) *eff = 0;
		if( !want ) return S_OK;
		size_t scn_num = _sections.size();
		size_t i = 0;
		for( ; i < scn_num; ++ i )
		{
			uint32_t this_va = _sections[i].VirtualAddress;
			uint32_t this_vs = _sections[i].Misc.VirtualSize;
			uint32_t this_ve = this_va + this_vs;
			if( this_ve < this_va ) 
			{
				if( i + 1 < scn_num ) 
					this_ve = _sections[i+1].VirtualAddress;
				else
					this_ve = _is64 ? _nt_hdr64.SizeOfImage : _nt_hdr32.SizeOfImage;

				if( this_ve < this_va )
					this_ve = -1;
			}

			if( rva >= this_va && rva < this_ve ) 
				break;
		}

		if( i == scn_num ) return 0;

		char * tofill = (char *)pbuf;
		size_t filded = 0;
		bool goon = true;
		for( ; goon && i < scn_num && filded < want; ++ i )
		{
			uint32_t this_va = _sections[i].VirtualAddress;
			uint32_t this_vs = _sections[i].Misc.VirtualSize;
			uint32_t this_ra = _sections[i].PointerToRawData;
			uint32_t this_rs = _sections[i].SizeOfRawData;

			uint32_t read_ofs = 0;
			if( this_va < rva )
			{
				read_ofs = (uint32_t)rva - this_va;
			}
			uint32_t read_vbs = this_vs - read_ofs;
			read_vbs = min( want - filded, read_vbs );
			memset( tofill, 0, read_vbs );
			if( this_rs && read_ofs < this_rs )
			{
				uint32_t r_pos = this_ra + read_ofs;
				uint32_t r_size = this_rs - read_ofs;
				r_size = min( r_size, read_vbs );
				size_t rbs = 0;
				long lr = _file->Readp( r_pos, tofill, r_size, &rbs );
				if( lr < 0 ) break;
				if( rbs < r_size ) goon = false;
			}
			tofill += read_vbs;
			filded += read_vbs;
		}
		if( eff ) *eff = filded;
		return S_OK;
	}
	STDMETHODIMP_(UINT64) imagex_t::Length()
	{
		if( _is64 ) return _nt_hdr64.SizeOfImage;
		return _nt_hdr32.SizeOfImage;
	};

	size_t imagex_t::num_of_segment()
	{
		return _sections.size();
	}
	long imagex_t::query_segment( size_t i, seginfo_t * seg )
	{ 
		if( i >= _sections.size() ) return E_INVALIDARG;
		IMAGE_SECTION_HEADER & scn = _sections[i];
		seg->offset = scn.VirtualAddress;
		seg->bytes = scn.Misc.VirtualSize;
		seg->paga = sc_to_paga( scn.Characteristics );
		return S_OK;
	}
	long imagex_t::find_segment( uint32_t rva, seginfo_t * vmip )
	{ 
		IMAGE_SECTION_HEADER * ph = find_section( rva, 0 );
		if( !ph ) return E_FAIL;
		vmip->offset = ph->VirtualAddress;
		vmip->bytes = ph->Misc.VirtualSize;
		vmip->paga = sc_to_paga( ph->Characteristics );
		return S_OK;
	}
	uint64_t imagex_t::get_image_base()
	{
		return image_base();
	}
	long imagex_t::relocate( vmarea_t * area, process_t * fixer )
	{
		int64_t diff = area->base - image_base();
		if( !diff ) return S_OK;

		IMAGE_DATA_DIRECTORY * pdde = get_dde( IMAGE_DIRECTORY_ENTRY_BASERELOC );
		if( !pdde || !pdde->VirtualAddress || !pdde->Size )
			return S_OK;

		size_t reloc_size = pdde->Size;
		size_t reloc_base = pdde->VirtualAddress;
		enum { MAX_RELOC_DATA_SIZE = 0x80000, };
		if( reloc_size > MAX_RELOC_DATA_SIZE ) reloc_size = MAX_RELOC_DATA_SIZE;

		UTIL::sentry<uint8_t*> reloc_page(new uint8_t[VM_PAGE_SIZE]);
		UTIL::sentry<uint8_t*> reloc_buf(new uint8_t[reloc_size]);
		RASSERT(reloc_buf && reloc_page, E_FAIL);

		size_t reloc_data_rest = 0;
		long lr = Readp( reloc_base, reloc_buf, reloc_size, &reloc_data_rest );
		if( lr < 0 ) return E_FAIL;

		RASSERT( reloc_data_rest >= sizeof(IMAGE_BASE_RELOCATION), S_OK );

		const uint8_t * reloc_data = reloc_buf.m_p;
		const uint8_t * reloc_data_end = reloc_data + reloc_data_rest;

		size_t eff = 0;
		for( ; reloc_data + sizeof(IMAGE_BASE_RELOCATION) <= reloc_data_end; )
		{
			IMAGE_BASE_RELOCATION * reloc_block = (IMAGE_BASE_RELOCATION *)(reloc_data);

			uint32_t reloc_page_addr = reloc_block->VirtualAddress;
			uint32_t reloc_ent_size = reloc_block->SizeOfBlock;

			if( reloc_ent_size < sizeof(IMAGE_BASE_RELOCATION) )
				reloc_ent_size = sizeof(IMAGE_BASE_RELOCATION);

			if( reloc_ent_size > reloc_data_rest )
				reloc_ent_size = reloc_data_rest;

			if( reloc_ent_size < sizeof(IMAGE_BASE_RELOCATION) )
				break;

			// move to next reloca-block
			reloc_data += reloc_ent_size;
			reloc_data_rest -= reloc_ent_size;

			// point to reloc-items
			const uint16_t * reloc_offset = (const uint16_t *)(reloc_block + 1);
			size_t reloc_offset_num = (reloc_ent_size - sizeof(IMAGE_BASE_RELOCATION))/sizeof(uint16_t);

			lr = Readp( reloc_page_addr, reloc_page, VM_PAGE_SIZE, &eff );
			if( lr < 0 ) break;
			seginfo_t seg;
			find_segment( reloc_page_addr, &seg );
			lr = fixer->relocate( area, reloc_page_addr, seg.paga, diff, reloc_offset, reloc_offset_num, reloc_page );
			if( lr < 0 ) break;
		}
		return lr;
	}

	long imagex_t::resolve_imports( vmarea_t * area, process_t * fixer )
	{
		IMAGE_DATA_DIRECTORY * pdde = get_dde( IMAGE_DIRECTORY_ENTRY_IMPORT );
		if( !pdde || !pdde->VirtualAddress || !pdde->Size )
			return S_OK;

		UTIL::sentry<char*> lib_name_buf = new char[MAX_PATH];
		if( !lib_name_buf ) return E_OUTOFMEMORY;

		long lr = 0;
		uint32_t imp_rva =  pdde->VirtualAddress;
		enum { MAX_LIBRARIES_TO_LAOD = 256, };
		for( size_t i = 0; i < MAX_LIBRARIES_TO_LAOD; ++ i, imp_rva += sizeof(IMAGE_IMPORT_DESCRIPTOR) )
		{
			IMAGE_IMPORT_DESCRIPTOR iidesc = {};
			size_t eff = 0;
			lr = Readp( imp_rva, &iidesc, sizeof(iidesc), &eff );
			if( lr < 0 ) break;
			if( eff < sizeof(iidesc) ) break;
			if( !iidesc.Characteristics ) break;
			if( !iidesc.Name ) break;
			lr = Readp( iidesc.Name, lib_name_buf, MAX_PATH, &eff );
			if( lr < 0 || eff < 2 ) continue;
			lib_name_buf[MAX_PATH-1] = 0;
			fix_ms_dllname( lib_name_buf );
			if( !lib_name_buf[0] ) continue;

			if( _is64 ) 
				resolve_thunks<IMAGE_THUNK_DATA64,uint64_t>( area, lib_name_buf, &iidesc, fixer );
			else
				resolve_thunks<IMAGE_THUNK_DATA32,uint32_t>( area, lib_name_buf, &iidesc, fixer );
		}
		return S_OK;
	};

	template< class IMAGE_THUNK_DATA3264, class PTRT >
	long imagex_t::resolve_thunks( vmarea_t * area, const char * libname, IMAGE_IMPORT_DESCRIPTOR * piidesc, process_t * fixer )
	{
		IMAGE_THUNK_DATA3264 thunk;

		const uint64_t IMP_BY_ORD_BIT = (1LL<<(sizeof(thunk)*8-1));

		UTIL::sentry<char*> api_name_buf = new char[MAX_PATH*2];
		if( !api_name_buf ) return E_OUTOFMEMORY;

		size_t eff = 0; long lr = 0;
		uint64_t thunk_rva = piidesc->OriginalFirstThunk;

		vmarea_t * imgarea;
		lr = fixer->image_load( &imgarea, libname, IMG_MAP_STATIC_LOAD|IMG_MAP_STD_LOAD );
		if( lr < 0 ) imgarea = 0;

		vxvma_t dll_base = imgarea ? imgarea->base : 0;
		optr<imagex_t> dll_img = imgarea ? imgarea->vmso->proto<image_view_t>()->image_object() : 0;

		for( size_t i = 0; i < MX_THUNKS_TO_PTOC; ++ i, thunk_rva += sizeof(thunk) )
		{
			lr = Readp( thunk_rva, &thunk, sizeof(thunk), &eff );
			if( lr < 0 ) break;
			if( eff != sizeof(thunk) ) break;
			if( !thunk.u1.AddressOfData ) break;

			uint64_t iat = piidesc->FirstThunk + sizeof(thunk)*i;

			uint32_t rva = 0;
			if( imgarea )
			{
				if( thunk.u1.Ordinal & IMP_BY_ORD_BIT )
				{
					rva = dll_img->get_proc_address_by_ord( (thunk.u1.Ordinal & 0xFFFF) );
				}
				else
				{
					uint64_t api_rva = thunk.u1.ForwarderString + sizeof(uint16_t);
					lr = Readp( api_rva, api_name_buf, MAX_PATH, &eff );
					if( lr < 0 ) continue;
					api_name_buf[MAX_PATH-1] = 0;
					rva = dll_img->get_proc_address_by_name( api_name_buf );
				}
			};
			PTRT api = rva ? rva + dll_base : fixer->gunkapi;
			uint64_t wc = sizeof(api);
			fixer->_vms_write( area, iat, &api, &wc );
		}
		return S_OK;
	}

	uint32_t imagex_t::get_proc_address_by_name( const char * name )
	{
		load_exports();
		name_ord_pair tmp; 
		tmp.name = name; tmp.fidx = 0;
		name_ord_vector::iterator it = std::lower_bound( 
			export_names.begin(), export_names.end(), tmp, less_name_ord_pair );
		if( it == export_names.end() ) 
			return 0;
		if( !(it->name == tmp.name) ) 
			return 0;
		size_t index = it->fidx;
		if( index >= exports.size() )
			return NULL;
		return exports[index];
	}

	uint32_t imagex_t::get_proc_address_by_idx( size_t idx )
	{
		load_exports();
		if( idx >= exports.size() ) return NULL;
		return exports[idx];
	}

	uint32_t imagex_t::get_proc_address_by_ord( size_t ord )
	{
		load_exports();
		if( ord < export_ord_base ) return NULL;
		size_t index = ord - export_ord_base;
		if( index >= exports.size() ) return NULL;
		return exports[index];
	}

	const char * imagex_t::get_proc_name_by_idx( size_t idx )
	{
		load_exports();
		if( idx >= exports.size() ) return NULL;

		o2n_t tmp; 
		tmp.fidx = idx; tmp.nidx = 0;

		std::vector<o2n_t>::iterator it = std::lower_bound( export_names_o2n.begin(), export_names_o2n.end(), tmp );

		if( it == export_names_o2n.end() ) 
			return NULL;

		if( it->fidx != tmp.fidx ) 
			return NULL;

		size_t nidx = it->nidx;
		if( nidx >= export_names.size() )
			return NULL;

		return export_names[nidx].name.c_str();
	}

	const char * imagex_t::enum_named_proc( size_t nidx, size_t * fidx )
	{
		load_exports();
		if( nidx >= export_names.size() )
			return NULL;
		if( fidx ) *fidx = export_names[nidx].fidx;
		return export_names[nidx].name.c_str();
	}


	long imagex_t::load_exports()
	{
		if( export_loads ) return S_OK;
		export_loads ++;

		IMAGE_DATA_DIRECTORY * dde = get_dde( IMAGE_DIRECTORY_ENTRY_EXPORT );
		if( !dde || !dde->VirtualAddress ||!dde->Size ) return S_FALSE;
		IMAGE_EXPORT_DIRECTORY dir_exp = {};
		size_t eff = 0;
		long lr = Readp( dde->VirtualAddress, &dir_exp, sizeof(dir_exp), &eff );
		if( lr < 0 || eff != sizeof(dir_exp) )
			return E_FAIL;

		const size_t MAX_EXPORTS = 0xFFFE;
		size_t func_num = dir_exp.NumberOfFunctions;
		size_t name_num = dir_exp.NumberOfNames;
		if( name_num > MAX_EXPORTS || func_num > MAX_EXPORTS ) 
			return E_FAIL;

		export_ord_base = dir_exp.Base;
		exports.resize( func_num );

		lr = Readp( dir_exp.AddressOfFunctions, exports.data(), sizeof(uint32_t)*func_num, &eff );
		if( lr < 0 ) return E_FAIL;

		func_num = eff/sizeof(uint32_t);
		exports.resize( func_num );

		UTIL::sentry<char*> api_name_buf = new char[MAX_PATH];
		if( !api_name_buf ) return E_OUTOFMEMORY;

		size_t to_read = name_num * sizeof(uint32_t);
		for( size_t offset = 0; offset < to_read; offset += sizeof(uint32_t) )
		{
			uint16_t name_to_ord;
			lr = Readp( dir_exp.AddressOfNameOrdinals + (offset>>1), &name_to_ord, sizeof(name_to_ord), &eff );
			if( lr < 0 ) 
				return E_FAIL;

			if( eff != sizeof(name_to_ord) ) 
				continue;

			if( name_to_ord >= func_num ) 
				continue;

			uint32_t name_rva;
			
			lr = Readp( dir_exp.AddressOfNames + offset, &name_rva, sizeof(uint32_t), &eff );
			
			if( lr < 0 ) 
				return E_FAIL;

			if( eff != sizeof(uint32_t) ) 
				continue;

			lr = Readp( name_rva, api_name_buf, MAX_PATH, &eff );
			if( lr < 0 || eff < 2 ) 
				continue;
			api_name_buf[MAX_PATH-1] = 0;
			if( !strlen(api_name_buf) ) 
				continue;

			name_ord_pair pair;
			pair.name = api_name_buf;
			pair.fidx = name_to_ord;
			export_names.push_back( pair );
		}

		if( export_names.size() )
		{
			std::sort( export_names.begin(), export_names.end(), less_name_ord_pair );

			for( size_t k = 0; k < export_names.size(); ++ k )
			{
				o2n_t o2n;
				o2n.fidx = export_names[k].fidx;
				o2n.nidx = k;
				export_names_o2n.push_back( o2n );
			}
			std::sort( export_names_o2n.begin(), export_names_o2n.end());
		}

		return S_OK;
	}

	size_t imagex_t::tls_rva()
	{
		IMAGE_DATA_DIRECTORY* dde = get_dde(IMAGE_DIRECTORY_ENTRY_TLS);
		if( !dde ) return NULL;
		return dde->VirtualAddress;
	}


};
